import numpy as np
import pandas as pd
from rdkit import Chem
from mordred import Calculator, descriptors
from rdkit.Chem import AllChem
from typing import Tuple, Dict, Generator

class MolecularFeatureEngine:
    def __init__(self):
        self.descriptor_calculator = Calculator(descriptors, ignore_3D=True)
        
    def sanitize_mol(self, smile: str) -> Chem.rdchem.Mol:
        """分子处理优化流水线"""
        try:
            mol = Chem.MolFromSmiles(smile)
            if not mol:
                return None
                
            mol = Chem.AddHs(mol)
            Chem.SanitizeMol(mol)
            AllChem.EmbedMolecule(mol)
            AllChem.UFFOptimizeMolecule(mol)
            return mol
        except:
            return None
    
    def batch_process_molecules(self, smiles: np.ndarray) -> Generator[Chem.rdchem.Mol, None, None]:
        """分子批量处理生成器"""
        for smile in smiles:
            mol = self.sanitize_mol(smile)
            if mol:
                yield mol
    
    def compute_descriptors(self, smiles_series: pd.Series) -> Tuple[pd.DataFrame, np.ndarray]:
        """改进版描述符计算方法"""
        mols = self.batch_process_molecules(smiles_series)
        
        # 预分配内存矩阵（优化内存管理）
        desc_values = self.descriptor_calculator.pandas([m for m in mols if m])
        desc_matrix = desc_values.astype(float).values
        
        # 改进的无效值处理
        desc_matrix = np.where(np.isinf(desc_matrix), np.nan, desc_matrix)
        desc_matrix = np.nan_to_num(desc_matrix, nan=np.nan)
        
        # 构建DataFrame时自动处理列名
        desc_df = pd.DataFrame(desc_matrix, 
                            columns=[str(d) for d in self.descriptor_calculator.descriptors])
        return desc_df, desc_df.columns.values

    def process_features(self, O_descs: pd.DataFrame, H_descs: pd.DataFrame) -> Dict[str, pd.DataFrame]:
        """内存矩阵特征工程"""
        # 类型转换优化
        O_matrix = O_descs.astype(np.float32).values
        H_matrix = H_descs.astype(np.float32).values
        
        # 特征差分计算
        diff_matrix = np.subtract(O_matrix, H_matrix, dtype=np.float32)
        
        # 条件调整计算
        mask = (diff_matrix >= 0) & (diff_matrix < 1e-6)
        adjusted_O = np.where(mask, O_matrix, diff_matrix)
        adjusted_H = np.where(mask, H_matrix, diff_matrix)
        
        # 构建完整结果集
        return {
            "O-H": pd.DataFrame(diff_matrix, columns=O_descs.columns),
        }

def main_pipeline(smiles_data: pd.DataFrame) -> pd.DataFrame:
    """主处理流水线"""
    engine = MolecularFeatureEngine()
    
    # 计算 O/H 描述符
    O_descs, desc_names = engine.compute_descriptors(smiles_data['O_smile'])
    H_descs, _ = engine.compute_descriptors(smiles_data['H_smile'])
    
    # 特征工程处理
    feature_dict = engine.process_features(O_descs, H_descs)
    
    # 合并最终结果
    final_df = pd.concat([
        smiles_data.reset_index(drop=True),
        feature_dict['O-H'].add_prefix('')
    ], axis=1)
    
    return final_df

# 使用示例
if __name__ == "__main__":
    # 原始数据加载
    raw_data = pd.read_csv("./test-data.csv")
    # filtered_data = raw_data[(raw_data['O_mol'] == True) & (raw_data['H_mol'] == True)]
    filtered_data = raw_data
    
    # 执行处理流水线
    result_df = main_pipeline(filtered_data)
    
    # 查看结果
    print(result_df.info())
    print(result_df["nBase"].head())